/*
 *  Copyright (C) 2010 Martin Haulrich <mwh.isv@cbs.dk>
 *
 *  This file is part of the MatrixParser package.
 *
 *  The MatrixParser program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Lesser General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Lesser General Public License for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package org.osdtsystem.matrixparser.parsers;

import org.osdtsystem.matrixparser.data.GraphIterator;
import org.osdtsystem.matrixparser.logging.Log;
import org.osdtsystem.matrixparser.parsers.loss.Loss;
import java.io.IOException;
import org.osdtsystem.matrixparser.data.CONLLIterator;
import org.osdtsystem.matrixparser.data.CONLLSentence;
import org.osdtsystem.matrixparser.parsers.output.DependencyTree;
import org.osdtsystem.matrixparser.features.FeatureVector;
import org.osdtsystem.matrixparser.learners.Learner;
import org.osdtsystem.matrixparser.features.WeightVector;
import org.osdtsystem.matrixparser.main.Options;

/**
 *
 * @author Martin Haulrich
 */
public class Trainer {
    final Learner learner;
    final Parser parser1, parser2;
    final Loss lossFunction;
    int sentences;

    public Trainer(Learner learner, Parser parser1, Loss lossFunction, Parser parser2) {
        this.learner = learner;
        this.parser1 = parser1;
        this.lossFunction = lossFunction;
        this.parser2 = parser2;
    }

    public Trainer(Learner learner, Parser trainParser, Loss loss) {
        this(learner, trainParser, loss, null);
    }

    public void train(CONLLIterator sentenceIterator,
            final int iterations1, final int iterations2)
            throws IOException {
        int tenPercent = Math.max(1, sentences / 10);
        long ta = System.currentTimeMillis();
        Log.info("Training..., " + sentences + " sentences, " + iterations1 + " MST iterations, " + iterations2 + " IP iterations");

        Parser parser = parser1;
        FeatureVector ignoredFeatures = learner.model().ignoredFeatures();
        if (iterations1 > 0)
            Log.config("Parser: ", parser.getClass());
        int[] iterationsList = {iterations1, iterations2};
        int parserID = 0;
        String modelFileName = Options.option("model.mst");
        for (int iterations : iterationsList) {
            parserID++;
            if (parserID == 2) {
                modelFileName = Options.option("model.irp");
                parser = parser2;
                Log.config("Parser: ", parser.getClass(), ignoredFeatures != null ? " (no updates of weights for MST features from this point)" : "");
            }
            for (int iter = 0; iter < iterations && parser != null; iter++) {
                Log.phandler.openIteration("Training iteration " + iter);
                double totalLoss = 0;
                long words = 0;
                sentenceIterator.reset();
                long a = System.currentTimeMillis();
                int sent = 0;
                while (sentenceIterator.hasNext()) {
                    CONLLSentence sentence = sentenceIterator.next();
                    Log.phandler.openSentence(sentence.toTokenString());
                    words += sentence.size() - 1;

                    // Parse sentence and create gold tree
                    WeightVector weights = learner.getWeightVector();
                    //System.err.println("weights[1]: " + learner.getWeightVector().checksum64());
                    DependencyTree systemTree = parser.parse(sentence, weights);
                    DependencyTree goldTree = parser.goldTree(sentence);
                    //System.err.println("weights[2]: " + learner.getWeightVector().checksum64());
                    if (Log.debugParserActions) {
                        System.err.println("gold:   " + goldTree.toString());
                        System.err.println("system: " + systemTree.toString());
                    }

                    // Compute feature vectors
                    FeatureVector goldFVs = parser.featureVector(parser.goldTree(sentence));
                    FeatureVector systemFVs = parser.featureVector(systemTree);
                    if (ignoredFeatures != null) {
                        goldFVs.ignoreFeatures(ignoredFeatures);
                        systemFVs.ignoreFeatures(ignoredFeatures);
                    }

                    // Compute loss
                    double loss = lossFunction.loss(goldTree, systemTree);
                    totalLoss += loss;

                    // Update weights
                    double update = (iterations - iter) * sentences - sent;
                    //double upd = iterations * numSen - (numSen * (iter) + (sent + 2) + 1);

                    learner.update(goldFVs, systemFVs,
                            parser.scorer(), loss, update);
                    //System.err.println("weights[3]: " + learner.getWeightVector().checksum64());
                    if ((sent + 1)  % tenPercent == 0) {
                        Log.infoW(".");
                    }
                    sent++;
                    Log.phandler.closeSentence();
                }

                if (iter == 0) {
                    System.gc();
                }
                long b = System.currentTimeMillis();
                long time = (b - a) / 1000;
                System.gc();
                Log.info("   finished iteration: " + (iter + 1) + "/"
                        + iterations + " [" + time + " seconds,  "
                        + learner.getWeightVector().size() + " weights, "
                        + totalLoss / words + " avg.loss/word, "
                        + ((int) ((Runtime.getRuntime().totalMemory()
                            - Runtime.getRuntime().freeMemory()) / 1024 / 1024)) + " MB heap]");
                //Log.setLevel(Level.ALL);
                Log.phandler.closeIteration();
            }
            //System.err.println("weights[4]: " + learner.getWeightVector().checksum64());
            learner.averageWeights();
        }
        //System.err.println("weights[5]: " + learner.getWeightVector().checksum64());
        long tb = System.currentTimeMillis();
        long ttime = (tb - ta) / 1000;
        Log.info("Finished training [" + ttime + " seconds]");
    }

    public void setSentences(int sentences) {
        this.sentences = sentences;
    }

    public void growAlphabet(GraphIterator trainIterator) {
        throw new UnsupportedOperationException("Not yet implemented");
    }

    public void train(GraphIterator trainIterator, GraphIterator devIterator) {
        throw new UnsupportedOperationException("Not yet implemented");
    }
}
